{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# 流式自定义生成器函数\n",
    "您可以在LCEL流水线中使用生成器函数（即使用`yield`关键字并像迭代器一样工作的函数）。\n",
    "\n",
    "这些生成器的签名应为`Iterator[Input] -> Iterator[Output]`。或者对于异步生成器：`AsyncIterator[Input] -> AsyncIterator[Output]`。\n",
    "\n",
    "这些对于以下情况非常有用：\n",
    "\n",
    "- 实现自定义输出解析器\n",
    "- 修改先前步骤的输出，同时保留流式处理能力\n",
    "\n",
    "让我们为逗号分隔列表实现一个自定义输出解析器\n",
    "\n",
    "## 同步版本\n"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "c77c7bb7d3a74cee"
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "|熊,| |狮子, |大象, |长颈鹿|, 企鹅||"
     ]
    },
    {
     "data": {
      "text/plain": "'熊, 狮子, 大象, 长颈鹿, 猴子'"
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from typing import Iterator, List\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "import os\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "template = \"\"\"\n",
    "    写一个逗号分隔的5只动物的列表: {animal}\n",
    "    \"\"\"\n",
    "\n",
    "prompt = ChatPromptTemplate.from_template(template)\n",
    "# 请注意，我们将max_retries = 0设置为避免在RateLimits等情况下重试\n",
    "llm = ChatOpenAI(\n",
    "    # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key=\"sk-xxx\",\n",
    "    openai_api_key=os.getenv(\"DASHSCOPE_API_KEY\"),\n",
    "    openai_api_base=\"https://dashscope.aliyuncs.com/compatible-mode/v1\",\n",
    "    model_name=\"qwen-max\",\n",
    "    temperature=0.0,\n",
    "    max_retries=0,\n",
    ")\n",
    "chain = prompt | llm | StrOutputParser()\n",
    "\n",
    "for chunk in chain.stream({\"animal\": \"熊\"}):\n",
    "    print(chunk, end=\"|\", flush=True)\n",
    "\n",
    "chain.invoke({\"animal\": \"熊\"})\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-10-28T08:45:02.068117Z",
     "start_time": "2024-10-28T08:44:57.670951Z"
    }
   },
   "id": "5f789fa2c9e68e8e",
   "execution_count": 15
  },
  {
   "cell_type": "markdown",
   "source": [
    "使用生成器Iterator[Input] -> Iterator[Output] 将llm令牌的迭代器拆分为逗号分隔的字符串列表"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "df64e0a4862f4605"
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<itertools._tee object at 0x111df1d00>\n",
      "['熊']\n",
      "['狮子']\n",
      "['大象']\n",
      "['猴子']\n",
      "['长颈鹿']\n"
     ]
    }
   ],
   "source": [
    "def split_into_list(input1: Iterator[str]) -> Iterator[List[str]]:\n",
    "    \"\"\"\n",
    "    将llm令牌的迭代器拆分为逗号分隔的字符串列表\n",
    "    :param input1: \n",
    "    :return: \n",
    "    \"\"\"\n",
    "    # 假设输入是逗号分隔的字符串\n",
    "    print(input1)\n",
    "    buffer = \"\"\n",
    "    for chunk1 in input1:\n",
    "        # 将chunk添加到缓冲区\n",
    "        buffer += chunk1\n",
    "        while \",\" in buffer:\n",
    "            # 在逗号处拆分缓冲区\n",
    "            comma_index = buffer.index(\",\")\n",
    "            # 在逗号之前的所有内容\n",
    "            yield [buffer[:comma_index].strip()]\n",
    "            # 保存剩余内容以供下一次迭代\n",
    "            buffer = buffer[comma_index + 1:]\n",
    "\n",
    "    # 返回剩余内容\n",
    "    yield [buffer.strip()]\n",
    "\n",
    "\n",
    "list_chain = chain | split_into_list\n",
    "\n",
    "for chunk2 in list_chain.stream({\"animal\": \"熊\"}):\n",
    "    print(chunk2, flush=True)\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-10-28T08:43:37.373817Z",
     "start_time": "2024-10-28T08:43:35.031919Z"
    }
   },
   "id": "1549eaf593245720",
   "execution_count": 14
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 异步版本\n"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "94d32b13f685d25d"
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['熊']\n",
      "['狮子']\n",
      "['大象']\n",
      "['长颈鹿']\n",
      "['猴子']\n"
     ]
    }
   ],
   "source": [
    "from typing import AsyncIterator\n",
    "\n",
    "\n",
    "async def asplit_into_list(input3: AsyncIterator[str]) -> AsyncIterator[List[str]]:\n",
    "    buffer = \"\"\n",
    "    async for chunk3 in input3:\n",
    "        buffer += chunk3\n",
    "        while \",\" in buffer:\n",
    "            comma_index = buffer.index(\",\")\n",
    "            yield [buffer[:comma_index].strip()]\n",
    "            buffer = buffer[comma_index + 1:]\n",
    "\n",
    "    yield [buffer.strip()]\n",
    "\n",
    "\n",
    "list_achain = chain | asplit_into_list\n",
    "\n",
    "async for chunk3 in list_achain.astream({\"animal\": \"熊\"}):\n",
    "    print(chunk3, flush=True)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-10-28T08:49:28.130084Z",
     "start_time": "2024-10-28T08:49:26.779884Z"
    }
   },
   "id": "dad65073cd9610c3",
   "execution_count": 17
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   },
   "id": "d058f0f8a4a0d2a8"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
